#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
'''
@File         :att_control_offb.py
@Description  :姿态控制，主要用于从机的编队控制
@Time         :2024/05/30 16:54:31
@Author       :Lin Yuheng
@Version      :1.0
'''

import rospy
from mavros_msgs.msg import AttitudeTarget, PositionTarget, State
from mavros_msgs.srv import CommandBool, SetMode
from geometry_msgs.msg import Pose, PoseStamped, Twist, TwistStamped
from nav_msgs.msg import Odometry
from fixedwing_sitl.msg import *

import numpy as np

from tools.formation_structure import *
from tools.mathlib import *
from tools.vector import *
from tools.filter import *
from tools.increment_pid_controller import *
from tools.params import *
from controller.vir_leader_formation_controller_test import VirLeaderFormationControllerTest

STATE_UAV = 101
CMD_UAV = 102
ARM = 1
TAKEOFF = 2
LAND = 3
POS = 4
VEL = 5
ACC = 6
POS_VEL = 7
POS_ACC = 8
VEL_ACC = 9
POS_VEL_ACC = 10
ATT = 11

class Controller:
    def __init__(self):

        self.planeID = 1
        self.uavID = "plane_1"
        self.leaderID = "plane_0"

        self.begin_time = rospy.Time()
        self.current_time = 0

        self.fwstates = FWstates()
        self.fw4cmd =  FWcmd()
        self.leaderstates =  Leaderstates()
        self.formtion_control_states = Formation_control_states()
        self.fwmonitor_flag = Fwmonitor()
        self.fw_cmd_mode = Fw_cmd_mode()
        self.fw_current_mode = Fw_current_mode()
        
        self.formation_controller = VirLeaderFormationControllerTest() # 编队控制器
        self.fw_control_mode_current = "MANUAL" # 当前控制模式
        self.fw_control_mode_last = "MANUAL" # 上一时刻控制模式
        
        self._cmd_type = POS

        self._fw_params = S_FwModelParams() # 飞机模型参数
        self._tecs_params = TECSParams()    # TECS控制器参数
        self._lateral_params = LateralControllerParams() # 横侧向控制器参数

        self._leader_states = S_FwLeaderStates()  # 领机信息
        self._thisfw_states = S_FwStates()        # 本机信息
        self._formation_cmd = S_FW4Cmd()          # 四通道控制指令
        self._formation_error = S_FwError()       # 编队误差以及偏差
        self._mix_Xerr_params = S_MixXerrParams() # 编队控制器混合X方向误差产生参数，编队控制器参数
        self._mix_Yerr_params = S_MixYerrParams() # 编队控制器混合Y方向误差产生参数，编队控制器参数
        self._formation_sp = S_FwSetPoint()       # 发布编队控制器控制状态

        # 最终直接发布给mavros的姿态环控制量
        # self.angle = np.zeros(3)
        # self.quat = np.zeros(4)
        self.angle = [0,0,0]
        self.quat = [0,0,0,0]
        self.target_raw_att = AttitudeTarget()
        
        # 位置环控制量
        self.target_raw_pose = PoseStamped()
        
        # 飞机mavros状态
        self.vehicle_state = State()
        self.should_exit = False

        ###### 1. 订阅和发布的话题 ######
        # 【订阅】本机所有的状态量
        self.fw_states_sub = rospy.Subscriber(
            self.uavID + "/fixed_wing_control/fw_state",
            FWstates,
            callback=self.fw_states_callback
        )

        # 【订阅】领机所有的状态量
        self.leader_states_sub = rospy.Subscriber(
            self.leaderID + "/fixed_wing_control/leader_state",
            Leaderstates,
            callback=self.leader_states_callback
        )

        # pub
        # 向fixed_wing_control/fw_cmd话题发送四通到控制量
        self.fw_cmd_pub = rospy.Publisher(
            self.uavID + "/fixed_wing_control/fw_cmd",
            FWcmd,
            queue_size=1,
        )
        
        # 向mavros发送位置环期望位置点
        self.local_sp_pos_pub = rospy.Publisher(
            self.uavID + "/mavros/setpoint_position/local",
            PoseStamped,
            queue_size=10
        )
        
        # 向mavros直接发送固定翼四通道控制量
        self.local_sp_att_pub = rospy.Publisher(
            self.uavID + "/mavros/setpoint_raw/attitude",
            AttitudeTarget,
            queue_size=10,
        )
        
        # 向fixed_wing_control/formation_error话题发送编队误差
        self.formation_error_pub = rospy.Publisher(
            self.uavID + "/fixed_wing_control/formation_error",
            Formation_control_states,
            queue_size=1,
        )
        
        self.mavros_sub = rospy.Subscriber(
            self.uavID + "/mavros/state",
            State,
            callback=self.mavros_state_callback,
        )
        
    def mavros_state_callback(self, msg: State):
        if self.vehicle_state.mode == State.MODE_PX4_OFFBOARD:
            if msg.mode != State.MODE_PX4_OFFBOARD:
                rospy.logwarn("Vehicle quit offboard mode, quitting!")
                self.should_exit = True
        self.vehicle_state = msg

    def fw_states_callback(self, msg:FWstates):
        self.fwstates = msg
        
    def leader_states_callback(self, msg:Leaderstates):
        self.leaderstates = msg

    def get_ros_time(self, begin:rospy.Time):
        time_now = rospy.Time.now()
        currTimeSec = time_now.secs - begin.secs
        currTimenSec = time_now.nsecs / 1e9 - begin.nsecs / 1e9
        return (currTimeSec + currTimenSec)

    def formation_states_pub(self):
        # TODO 完成编队信息发布
        # 编队误差信息发布
        formation_control_states = Formation_control_states()
        formation_control_states.err_P_N = self._formation_error.P_N
        formation_control_states.err_P_E = self._formation_error.P_E
        formation_control_states.err_P_D = self._formation_error.P_D
        formation_control_states.err_P_NE = self._formation_error.P_NE
        
        formation_control_states.err_PXb = self._formation_error.PXb
        formation_control_states.err_PYb = self._formation_error.PYb
        formation_control_states.err_PZb = self._formation_error.PZb
        
        formation_control_states.err_PXk = self._formation_error.PXk
        formation_control_states.err_PYk = self._formation_error.PYk
        formation_control_states.err_PZk = self._formation_error.PZk
        
        formation_control_states.err_VXb = self._formation_error.VXb
        formation_control_states.err_VYb = self._formation_error.VYb
        formation_control_states.err_VZb = self._formation_error.VZb
        
        formation_control_states.err_VXk = self._formation_error.VXk
        formation_control_states.err_VYk = self._formation_error.VYk
        formation_control_states.err_VZk = self._formation_error.VZk
        
        formation_control_states.led_fol_vxb = self._formation_error.led_fol_vxb
        formation_control_states.led_fol_vyb = self._formation_error.led_fol_vyb
        formation_control_states.led_fol_vzb = self._formation_error.led_fol_vzb
        
        formation_control_states.led_fol_vxk = self._formation_error.led_fol_vxk
        formation_control_states.led_fol_vyk = self._formation_error.led_fol_vyk
        formation_control_states.led_fol_vzk = self._formation_error.led_fol_vzk
        
        formation_control_states.led_fol_eta = self._formation_error.led_fol_eta
        
        self.formation_error_pub.publish(formation_control_states)

    def control_formation(self):
        """
        @Description : 编队控制器主函数，完成对于领机从机状态的赋值，传入编队控制器
        @Param:
        @Returns     :
        """
        
        self.fw_control_mode_current = self.fwstates.control_mode
        
        self._cmd_type = self.leaderstates.cmd_type

        # 领机状态赋值,把从回调函数中获得的状态信息赋值给状态变量
        self._leader_states.air_speed = self.leaderstates.airspeed

        self._leader_states.altitude = self.leaderstates.altitude
        self._leader_states.latitude = self.leaderstates.latitude
        self._leader_states.longitude = self.leaderstates.longitude

        self._leader_states.ned_pos_n = self.leaderstates.ned_pos_x
        self._leader_states.ned_pos_e = self.leaderstates.ned_pos_y
        self._leader_states.ned_pos_d = self.leaderstates.ned_pos_z
        self._leader_states.relative_alt = self.leaderstates.relative_alt

        self._leader_states.ned_vel_n = self.leaderstates.ned_vel_x
        self._leader_states.ned_vel_e = self.leaderstates.ned_vel_y
        self._leader_states.ned_vel_d = self.leaderstates.ned_vel_z

        self._leader_states.global_vel_n = self.leaderstates.global_vel_x
        self._leader_states.global_vel_e = self.leaderstates.global_vel_y
        self._leader_states.global_vel_d = self.leaderstates.global_vel_z

        self._leader_states.roll_angle = self.leaderstates.roll_angle
        self._leader_states.pitch_angle = self.leaderstates.pitch_angle
        self._leader_states.yaw_angle = self.leaderstates.yaw_angle
        self._leader_states.yaw_valid = False # 领机的yaw不能直接获得

        # 从机状态赋值
        self._thisfw_states.flight_mode = self.fwstates.control_mode
        
        self._thisfw_states.air_speed = self.fwstates.air_speed 
        self._thisfw_states.in_air = self.fwstates.in_air

        self._thisfw_states.altitude = self.fwstates.altitude
        self._thisfw_states.altitude_lock = True # 保证TECS
        self._thisfw_states.in_air = True        # 保证TECS
        self._thisfw_states.latitude = self.fwstates.latitude
        self._thisfw_states.longitude = self.fwstates.longitude
        self._thisfw_states.ned_pos_n = self.fwstates.ned_pos_x
        self._thisfw_states.ned_pos_e = self.fwstates.ned_pos_y
        self._thisfw_states.ned_pos_d = self.fwstates.ned_pos_z
        
        self._thisfw_states.relative_alt = self.fwstates.relative_alt

        self._thisfw_states.ned_vel_n = self.fwstates.ned_vel_x
        self._thisfw_states.ned_vel_e = self.fwstates.ned_vel_y
        self._thisfw_states.ned_vel_d = self.fwstates.ned_vel_z

        self._thisfw_states.global_vel_n = self.fwstates.global_vel_x
        self._thisfw_states.global_vel_e = self.fwstates.global_vel_y
        self._thisfw_states.global_vel_d = self.fwstates.global_vel_z

        self._thisfw_states.roll_angle = self.fwstates.roll_angle
        self._thisfw_states.pitch_angle = self.fwstates.pitch_angle
        self._thisfw_states.yaw_angle = self.fwstates.yaw_angle
        self._thisfw_states.att_quat[0] = self.fwstates.att_quater.w
        self._thisfw_states.att_quat[1] = self.fwstates.att_quater.x
        self._thisfw_states.att_quat[2] = self.fwstates.att_quater.y
        self._thisfw_states.att_quat[3] = self.fwstates.att_quater.z
        self._thisfw_states.rotmat = quat_2_rotmax(self._thisfw_states.att_quat)

        self._thisfw_states.body_acc[0] = self.fwstates.body_acc_x
        self._thisfw_states.body_acc[1] = self.fwstates.body_acc_y
        self._thisfw_states.body_acc[2] = self.fwstates.body_acc_z
        self._thisfw_states.ned_acc = matrix_product_vector_3(self._thisfw_states.rotmat, self._thisfw_states.body_acc)

        self._thisfw_states.wind_estimate_x = self.fwstates.wind_estimate_x
        self._thisfw_states.wind_estimate_y = self.fwstates.wind_estimate_y
        self._thisfw_states.wind_estimate_z = self.fwstates.wind_estimate_z 

        # 设定编队形状 # 
        # self.formation_controller.set_formation_type()

        # 设定前向编队混合参数
        self.formation_controller.set_mix_Xerr_params(self._mix_Xerr_params)

        # 设定侧向编队混合参数
        self.formation_controller.set_mix_Yerr_params(self._mix_Yerr_params)

        # 如果模式不一致,需要重置控制器,达到控制连续
        if self.fw_control_mode_current != self.fw_control_mode_last:
            self.formation_controller.reset_formation_controller()
        
        # 更新飞机状态
        self.formation_controller.update_led_fol_states(self._leader_states, self._thisfw_states)

        # 编队控制
        self.formation_controller.control_formation()

        # 获得最终控制量
        self._formation_cmd = self.formation_controller.get_formation_4cmd()
        # 获得编队控制期望
        self._formation_sp = self.formation_controller.get_formation_sp()
        # 获得编队控制误差
        self._formation_error = self.formation_controller.get_formation_error()

        # 控制量赋值
        self.fw4cmd.throttle_sp = self._formation_cmd.throttle
        self.fw4cmd.roll_angle_sp = self._formation_cmd.roll
        self.fw4cmd.pitch_angle_sp = self._formation_cmd.pitch
        self.fw4cmd.yaw_angle_sp = self._formation_cmd.yaw

        # 发布四通道控制量
        self.fw_cmd_pub.publish(self.fw4cmd)

        #TODO 发布编队控制器状态
        self.formation_states_pub()
        
        self.fw_control_mode_last = self.fw_control_mode_current # 更新控制模式

    def start(self):
        print("start fixed-wing attitude control")
        rospy.init_node(
            self.uavID + "_attitude_control_node"
        )
        rate = rospy.Rate(20)

        while not rospy.is_shutdown():
            if self.should_exit:
                break
            try:
                self.control_formation()
                
                if self._cmd_type == POS:
                    rospy.loginfo("---位置环控制模式---")
                    # target_raw_pose 为 enu坐标系
                    self.target_raw_pose.pose.position.x = self._leader_states.ned_pos_e
                    self.target_raw_pose.pose.position.y = self._leader_states.ned_pos_n
                    self.target_raw_pose.pose.position.z = abs(self._leader_states.ned_pos_d)
                    
                    self.local_sp_pos_pub.publish(self.target_raw_pose)
                
                elif self._cmd_type == ATT:
                    rospy.loginfo("---姿态环控制模式---")
            
                    self.angle[0] = self.fw4cmd.roll_angle_sp
                    self.angle[1] = - self.fw4cmd.pitch_angle_sp
                    self.angle[2] = - self.fw4cmd.yaw_angle_sp + deg_2_rad(90.0)

                    # print("期望欧拉角： row:",self.angle[0], " pitch:",self.angle[1], " yaw:", self.angle[2], "throttle:", self.fw4cmd.throttle_sp)

                    self.quat = euler_2_quat(self.angle)

                    self.target_raw_att.type_mask = 7 # 1+2+4+64+128 body.rate_x,body.rate_y,body.rate_z thrust..
                    self.target_raw_att.orientation.w = self.quat[0]
                    self.target_raw_att.orientation.x = self.quat[1]
                    self.target_raw_att.orientation.y = self.quat[2]
                    self.target_raw_att.orientation.z = self.quat[3]
                    self.target_raw_att.thrust = self.fw4cmd.throttle_sp

                    self.local_sp_att_pub.publish(self.target_raw_att)
                    
                else:
                    rospy.loginfo("控制模式有误！")
                    break
            
                rate.sleep()
                
            except rospy.ROSException as e:
                pass
        
        rospy.loginfo("offb control end")

if __name__ == "__main__":
    Con = Controller()
    Con.start()
